package cn.com.doone.common.uc.infrastructure.shiro;

import java.util.Collection;
import java.util.LinkedHashSet;
import java.util.Set;

import javax.naming.NamingEnumeration;
import javax.naming.NamingException;
import javax.naming.directory.Attribute;
import javax.naming.directory.Attributes;
import javax.naming.directory.SearchControls;
import javax.naming.directory.SearchResult;
import javax.naming.ldap.LdapContext;

import org.apache.shiro.authc.AuthenticationInfo;
import org.apache.shiro.authc.AuthenticationToken;
import org.apache.shiro.authc.UsernamePasswordToken;
import org.apache.shiro.realm.activedirectory.ActiveDirectoryRealm;
import org.apache.shiro.realm.ldap.LdapContextFactory;
import org.apache.shiro.realm.ldap.LdapUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;

public class DooneActiveDirectoryRealm extends ActiveDirectoryRealm implements InitializingBean {

    private static final Logger log = LoggerFactory.getLogger(DooneActiveDirectoryRealm.class);

    @Override
    protected AuthenticationInfo queryForAuthenticationInfo(AuthenticationToken token, LdapContextFactory ldapContextFactory) throws NamingException {
        //final AuthenticationInfo queryForAuthenticationInfo = super.queryForAuthenticationInfo(token, ldapContextFactory);
        final UsernamePasswordToken upToken = (UsernamePasswordToken) token;
        LdapContext ctx = null;
        LdapContext passwordCtx = null;
        try {
            ctx = ldapContextFactory.getSystemLdapContext(); // .getLdapContext(upToken.getUsername(), upToken.getPassword());
            final String attribName = "sAMAccountName";
            final SearchControls searchCtls = new SearchControls(SearchControls.SUBTREE_SCOPE,1,0,new String[]{attribName, "distinguishedName"},false,false); 
            final NamingEnumeration<SearchResult> search = ctx.search(searchBase, "sAMAccountName={0}", new Object[]{upToken.getPrincipal()}, searchCtls); 
            // final NamingEnumeration<SearchResult> search = ctx.search(searchBase, "sAMAccountName=" + upToken.getUsername() , searchCtls);
            if(search.hasMore()) {
                final SearchResult next = search.next();
                upToken.setUsername(next.getAttributes().get("sAMAccountName").get().toString()); 
                passwordCtx = ldapContextFactory.getLdapContext(next.getAttributes().get("distinguishedName").get().toString(), upToken.getPassword());
            }
        } catch (NamingException ne) {
            log.error("Error in ldap name resolving", ne);
            // We have to rethrow the exception, to indicate invalid login
            throw ne;
        } finally {
            LdapUtils.closeContext(ctx);
            LdapUtils.closeContext(passwordCtx);
        }
        return buildAuthenticationInfo(upToken.getUsername(), upToken.getPassword());
    }

    private Set<String> getRoleNamesForUser(String username, LdapContext ldapContext) throws NamingException { 
        Set<String> roleNames; 
        roleNames = new LinkedHashSet<String>(); 
 
        SearchControls searchCtls = new SearchControls(); 
        searchCtls.setSearchScope(SearchControls.SUBTREE_SCOPE); 
 
        /*String userPrincipalName = username; 
        if (principalSuffix != null) { 
            userPrincipalName += principalSuffix; 
        }*/ 
 
        //SHIRO-115 - prevent potential code injection: 
        String searchFilter = "(&(objectClass=*)(sAMAccountName={0}))"; 
        Object[] searchArguments = new Object[]{username}; 
 
        NamingEnumeration answer = ldapContext.search(searchBase, searchFilter, searchArguments, searchCtls); 
 
        while (answer.hasMoreElements()) { 
            SearchResult sr = (SearchResult) answer.next(); 
 
            if (log.isDebugEnabled()) { 
                log.debug("Retrieving group names for user [" + sr.getName() + "]"); 
            } 
 
            Attributes attrs = sr.getAttributes(); 
 
            if (attrs != null) { 
                NamingEnumeration ae = attrs.getAll(); 
                while (ae.hasMore()) { 
                    Attribute attr = (Attribute) ae.next(); 
 
                    if (attr.getID().equals("memberOf")) { 
 
                        Collection<String> groupNames = LdapUtils.getAllAttributeValues(attr); 
 
                        if (log.isDebugEnabled()) { 
                            log.debug("Groups found for user [" + username + "]: " + groupNames); 
                        } 
 
                        Collection<String> rolesForGroups = getRoleNamesForGroups(groupNames); 
                        roleNames.addAll(rolesForGroups); 
                    } 
                } 
            } 
        } 
        return roleNames; 
    } 

	@Override
	public void afterPropertiesSet() throws Exception {
		// TODO Auto-generated method stub
	}

}